#pragma once

#include "Resource.h"

#include "ShaderCompiler.h"

#include "DynamicRHI.h"

namespace RenderSystem
{
	// Forward declarations.
	class VertexFactoryType;
	class ShaderType;

	class MaterialShaderType;
	class GlobalShaderType;
	class MeshMaterialShaderType;

	/** A shader parameter's register binding. */
	class ShaderParameter
	{
	public:
		ShaderParameter(): m_NumBytes(0) {}

		void Bind(const ShaderParameterMap& rParameterMap,const TCHAR* pszParameterName, bool bOptional = false );
		
		UINT GetBufferIndex() const { return m_BufferIndex; }
		UINT GetBaseIndex() const { return m_BaseIndex; }
		UINT GetNumBytes() const { return m_NumBytes; }
	
	private:
		WORD m_BufferIndex;
		WORD m_BaseIndex;
		WORD m_NumBytes;
	};

	/** A shader texture's sampler binding. */
	class ShaderTextureParameter
	{
	public:
		ShaderTextureParameter(): m_NumTextures(0) {}

		void Bind(const ShaderParameterMap& ParameterMap,const TCHAR* ParameterName, bool bOptional = false );
		
		UINT GetBaseIndex() const { return m_BaseIndex; }
		UINT GetNumTextures() const { return m_NumTextures; }
		UINT GetSamplerIndex() const { return m_SamplerIndex; }

	private:
		WORD m_BaseIndex;
		WORD m_NumTextures;
		WORD m_SamplerIndex;
	};

	/**
	* A compiled shader and its parameter bindings.
	*/
	class ShaderKey
	{
	public:

		ShaderKey() : 
		  m_dwParameterMapCRC(0)
		  {}

		  ShaderKey(const String& rStrInCode, const ShaderParameterMap& rInParameterMap) : 
		  m_strShaderCode(rStrInCode)
		  {
			  m_dwParameterMapCRC = rInParameterMap.GetCRC();
		  }

		  String	m_strShaderCode;
		  DWORD		m_dwParameterMapCRC;
	};

	class Shader
	{
	public:

		struct CompiledShaderInitializerType
		{
			ShaderType*					Type;
			
			String						strCodes;
			const ShaderParameterMap&	ParameterMap;
			
			UINT						NumInstructions;

			CompiledShaderInitializerType(
				ShaderType* InType,
				const ShaderCompilerOutput& rCompilerOutput
				):
				Type(InType),
				strCodes( rCompilerOutput.m_strCodes ),
				ParameterMap( rCompilerOutput.m_ParameterMap ),
				NumInstructions( rCompilerOutput.m_NumInstructions )
			{}
		};

		Shader( void );
		Shader( const CompiledShaderInitializerType& rInitializer );
		virtual ~Shader();

		/**
		* Can be overridden by Shader subclasses to modify their compile environment just before compilation occurs.
		*/
		static void ModifyCompilationEnvironment( ShaderCompilerEnvironment& OutEnvironment ) {}

		/**
		* @return the shader's vertex shader
		*/
		const FVertexShaderRHIRef& GetVertexShader(void) { return m_pVertexShader; }
		/**
		* @return the shader's vertex shader
		*/
		const FPixelShaderRHIRef& GetPixelShader(void) { return m_pPixelShader; }

		// Accessors.
		ShaderType* GetType() const { return m_pShaderType; }

		const ShaderKey& GetShaderKey(void) const { return m_ShaderKey; }

	protected:

		ShaderType*				m_pShaderType;

		FVertexShaderRHIRef		m_pVertexShader;
		FPixelShaderRHIRef		m_pPixelShader;

		ShaderKey				m_ShaderKey;

		UINT					m_uNumInstructions;
	};

	/**
	* An object which is used to serialize/deserialize, compile, and cache a particular shader class.
	*/
	class ShaderType
	{
	public:

		typedef Shader* (*ConstructSerializedType)();
		typedef void (*ModifyCompilationEnvironmentType)(ShaderCompilerEnvironment&);

		typedef std::map< uint32, Shader* >		IdShaders;
		typedef IdShaders::iterator				IdShadersIterator;
		typedef IdShaders::const_iterator		IdShadersConstIterator;

		typedef std::set< ShaderType* >			ShaderTypes;
		typedef ShaderTypes::iterator			ShaderTypesIterator;
		typedef ShaderTypes::const_iterator		ShaderTypesConstIterator;

		/**
		* Minimal initialization constructor.
		*/
		ShaderType(
			const TCHAR* InName,
			const TCHAR* InSourceFilename,
			const TCHAR* InFunctionName,
			ConstructSerializedType InConstructSerializedRef,
			ModifyCompilationEnvironmentType InModifyCompilationEnvironmentRef
		):
			m_szShaderTypeName(InName),
			m_szSourceFilename(InSourceFilename),
			m_szEntryFunctionName(InFunctionName),
			ConstructSerializedRef(InConstructSerializedRef),
			ModifyCompilationEnvironmentRef(InModifyCompilationEnvironmentRef)
		{
		}

		/**
		* Registers a shader for lookup by ID or code.
		*/
		void RegisterShader(Shader* Shader);

		/**
		* Removes a shader from the ID and code lookup maps.
		*/
		void UnregisterShader(Shader* Shader);

		/**
		* Constructs a new instance of the shader type for deserialization.
		*/
		Shader* ConstructForDeserialization() const;

		const TCHAR* GetName() const 
		{ 
			return m_szShaderTypeName; 
		}
		const TCHAR* GetShaderFilename() const 
		{ 
			return m_szSourceFilename; 
		}

		static ShaderTypes& GetTypeSets();

		bool HasShader() { return !m_IdShaders.empty(); }

	public:

		virtual MaterialShaderType* GetMaterialShaderType() { return NULL; }
		virtual GlobalShaderType* GetGlobalShaderType() { return NULL; }
		virtual MeshMaterialShaderType* GetMeshMaterialShaderType() { return NULL; }

	protected:

		/**
		* Enqueues a shader to be compiled with the shader type's compilation parameters, using the provided shader environment.
		* @param VFType - Optional vertex factory type that the shader belongs to.
		* @param Platform - Platform to compile for.
		* @param Environment - The environment to compile the shader in.
		*/
		void BeginCompileShader(VertexFactoryType* pVFType, const ShaderProfile& rShaderProfile, const ShaderCompilerEnvironment& rEnvironment);

		Shader* _FindShaderByOutput( const ShaderCompilerOutput& rOutput );

	protected:
		const TCHAR*						m_szShaderTypeName;
		const TCHAR*						m_szSourceFilename;
		const TCHAR*						m_szEntryFunctionName;
		IdShaders							m_IdShaders;

		
		ConstructSerializedType				ConstructSerializedRef;
		ModifyCompilationEnvironmentType	ModifyCompilationEnvironmentRef;

		class ShaderCodeKeyFuncs
			:public std::binary_function<Shader*, Shader*, bool>
		{
			bool operator() ( const Shader* pShaderLHS, const Shader* pShaderRHS ) const
			{
				const ShaderKey& key1 = pShaderLHS->GetShaderKey();
				const ShaderKey& key2 = pShaderRHS->GetShaderKey();

				return key1.m_dwParameterMapCRC == key2.m_dwParameterMapCRC;
			}
		};

		/** A map from shader code to shader. */
		typedef std::set< Shader*, ShaderCodeKeyFuncs > ShaderCodeMaps;
		typedef ShaderCodeMaps::iterator				ShaderCodeMapsIterator;
		typedef ShaderCodeMaps::const_iterator			ShaderCodeMapsConstIterator;

		ShaderCodeMaps			m_ShaderCodeMap;

	};

	/**
	* A macro to declare a new shader type.  This should be called in the class body of the new shader type.
	* @param ShaderClass - The name of the class representing an instance of the shader type.
	* @param ShaderMetaTypeShortcut - The shortcut for the shader meta type: simple, material, meshmaterial, etc.  The shader meta type
	*	controls 
	*/
	#define DECLARE_SHADER_TYPE(ShaderClass,ShaderMetaTypeShortcut) \
	public: \
		typedef ShaderMetaTypeShortcut##ShaderType ShaderMetaType; \
		static ShaderMetaType StaticType; \
		static Shader* ConstructSerializedInstance() { return new ShaderClass(); } \
		static Shader* ConstructCompiledInstance(const ShaderMetaType::CompiledShaderInitializerType& Initializer) \
		{ return new ShaderClass(Initializer); }

	/**
	* A macro to implement a shader type.
	*/
	#define IMPLEMENT_SHADER_TYPE(TemplatePrefix,ShaderClass,SourceFilename,FunctionName) \
		TemplatePrefix \
		ShaderClass::ShaderMetaType ShaderClass::StaticType( \
		TEXT(#ShaderClass), \
		SourceFilename, \
		FunctionName, \
		ShaderClass::ConstructSerializedInstance, \
		ShaderClass::ConstructCompiledInstance, \
		ShaderClass::ModifyCompilationEnvironment \
		);

	/**
	* A collection of shaders of different types, but the same meta type.
	*/
	template<typename ShaderMetaType>
	class TShaderMap
	{
	public:

		typedef std::map< typename ShaderMetaType*, Shader* >		ShaderMaps;
		//typedef ShaderMaps::iterator						ShaderMapsIterator;
		//typedef ShaderMaps::const_iterator					ShaderMapsConstIterator;

		TShaderMap(){}
		virtual ~TShaderMap(){}

		void AddShader( ShaderMetaType* pShaderType, Shader* pShader )
		{
			m_ShaderTypes.insert( ShaderMaps::value_type(pShaderType, pShader) );
		}

		/** Finds the shader with the given type. */
		bool HasShader(ShaderMetaType* pType) const
		{
			ShaderTypesConstIterator itrFound = m_ShaderTypes.find( pType );
			return itrFound != m_ShaderTypes.end();
		}

	private:
		ShaderMaps			m_ShaderTypes;
	};

}